1831F - Mex Tree - CodeForces Solution

#include <iostream>
#include <algorithm>
#include <vector>
#include <array>

using namespace std;
using ll = long long;

const int maxn = 2e5 + 10, S = 1000;

vector<vector<int>> g(maxn);
vector<int> sz(maxn), fa(maxn);

void dfs_sz(int u) {
    if (fa[u]) {
        for (auto it = g[u].begin(); it != g[u].end(); ++it) 
            if (*it == fa[u]) {
                g[u].erase(it);
                break;
            } 
    }

    sz[u] = 1;
    for (auto &v : g[u]) {
        fa[v] = u;
        dfs_sz(v);
        sz[u] += sz[v];
        if (sz[v] > sz[g[u][0]]) swap(v, g[u][0]);
    }
} 

const ll INF = 1e18;
struct Info {
    vector<array<ll, 2>> dp;
    Info(int n) : dp(n, {INF, INF}) {};
};
Info operator+ (const Info &a, const Info &b) {
    const int n = a.dp.size() - 1, m = b.dp.size() - 1;
    Info ans(min(n + m, S) + 1);
    for (int i = 1; i <= n; ++i)
        for (int j = 1; j <= m; ++j) {
            ans.dp[i][0] = min(ans.dp[i][0], a.dp[i][0] + b.dp[j][1]);
            ans.dp[i][1] = min(ans.dp[i][1], a.dp[i][1] + b.dp[j][0]);
            if (i + j <= S) {
                ans.dp[i + j][0] = min(ans.dp[i + j][0], a.dp[i][0] + b.dp[j][0] + i * j);
                ans.dp[i + j][1] = min(ans.dp[i + j][1], a.dp[i][1] + b.dp[j][1] + 2 * i * j);
            }
        }
    return ans;
}

Info dfs(int u) {
    Info ans(2);
    ans.dp[1][0] = 1;
    ans.dp[1][1] = 2;
    for (auto v : g[u]) {
        auto f = dfs(v);
        ans = ans + f;
    }
    return ans;
}

int main(){
	std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);

    int T;
    std::cin >> T;
    while (T--) {
        int n;
        std::cin >> n;
        for (int i = 1; i <= n; ++i) g[i].clear();
        for (int i = 1; i < n; ++i) {
            int u, v;
            std::cin >> u >> v;
            g[u].push_back(v), g[v].push_back(u);
        }
        dfs_sz(1);
        auto dp = dfs(1).dp;
        ll ans = INF;
        for (int i = 1; i <= min(S, sz[1]); ++i) {
            ans = min(ans, dp[i][0]);
            ans = min(ans, dp[i][1]);
        }
        std::cout << 1ll * n * (n + 1) - ans << '\n';
    }
    
	return 0;
}